/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.lucene.search;

import java.io.IOException;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.FilteredTermsEnum;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.MultiReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.analysis.MockAnalyzer;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.util.LuceneTestCase;
import org.apache.lucene.util.AttributeSource;
import org.apache.lucene.util.BytesRef;
import org.junit.AfterClass;
import org.junit.BeforeClass;

public class TestMultiTermQueryRewrites extends LuceneTestCase {

  static Directory dir, sdir1, sdir2;
  static IndexReader reader, multiReader, multiReaderDupls;
  static IndexSearcher searcher, multiSearcher, multiSearcherDupls;

  @BeforeClass
  public static void beforeClass() throws Exception {
    dir = newDirectory();
    sdir1 = newDirectory();
    sdir2 = newDirectory();
    final RandomIndexWriter writer =
        new RandomIndexWriter(random(), dir, new MockAnalyzer(random()));
    final RandomIndexWriter swriter1 =
        new RandomIndexWriter(random(), sdir1, new MockAnalyzer(random()));
    final RandomIndexWriter swriter2 =
        new RandomIndexWriter(random(), sdir2, new MockAnalyzer(random()));

    for (int i = 0; i < 10; i++) {
      Document doc = new Document();
      doc.add(newStringField("data", Integer.toString(i), Field.Store.NO));
      writer.addDocument(doc);
      ((i % 2 == 0) ? swriter1 : swriter2).addDocument(doc);
    }
    writer.forceMerge(1);
    swriter1.forceMerge(1);
    swriter2.forceMerge(1);
    writer.close();
    swriter1.close();
    swriter2.close();

    reader = DirectoryReader.open(dir);
    searcher = newSearcher(reader);

    multiReader =
        new MultiReader(
            new IndexReader[] {DirectoryReader.open(sdir1), DirectoryReader.open(sdir2)}, true);
    multiSearcher = newSearcher(multiReader);

    multiReaderDupls =
        new MultiReader(
            new IndexReader[] {DirectoryReader.open(sdir1), DirectoryReader.open(dir)}, true);
    multiSearcherDupls = newSearcher(multiReaderDupls);
  }

  @AfterClass
  public static void afterClass() throws Exception {
    reader.close();
    multiReader.close();
    multiReaderDupls.close();
    dir.close();
    sdir1.close();
    sdir2.close();
    reader = multiReader = multiReaderDupls = null;
    searcher = multiSearcher = multiSearcherDupls = null;
    dir = sdir1 = sdir2 = null;
  }

  private Query extractInnerQuery(Query q) {
    if (q instanceof ConstantScoreQuery) {
      // wrapped as ConstantScoreQuery
      q = ((ConstantScoreQuery) q).getQuery();
    }
    return q;
  }

  private Term extractTerm(Query q) {
    q = extractInnerQuery(q);
    return ((TermQuery) q).getTerm();
  }

  private void checkBooleanQueryOrder(Query q) {
    q = extractInnerQuery(q);
    final BooleanQuery bq = (BooleanQuery) q;
    Term last = null, act;
    for (BooleanClause clause : bq.clauses()) {
      act = extractTerm(clause.query());
      if (last != null) {
        assertTrue("sort order of terms in BQ violated", last.compareTo(act) < 0);
      }
      last = act;
    }
  }

  private void checkDuplicateTerms(MultiTermQuery.RewriteMethod method) throws Exception {
    final MultiTermQuery mtq = TermRangeQuery.newStringRange("data", "2", "7", true, true, method);
    final Query q1 = searcher.rewrite(mtq);
    final Query q2 = multiSearcher.rewrite(mtq);
    final Query q3 = multiSearcherDupls.rewrite(mtq);
    if (VERBOSE) {
      System.out.println();
      System.out.println("single segment: " + q1);
      System.out.println("multi segment: " + q2);
      System.out.println("multi segment with duplicates: " + q3);
    }
    assertEquals("The multi-segment case must produce same rewritten query", q1, q2);
    assertEquals(
        "The multi-segment case with duplicates must produce same rewritten query", q1, q3);
    checkBooleanQueryOrder(q1);
    checkBooleanQueryOrder(q2);
    checkBooleanQueryOrder(q3);
  }

  public void testRewritesWithDuplicateTerms() throws Exception {
    checkDuplicateTerms(MultiTermQuery.SCORING_BOOLEAN_REWRITE);

    checkDuplicateTerms(MultiTermQuery.CONSTANT_SCORE_BOOLEAN_REWRITE);

    // use a large PQ here to only test duplicate terms and dont mix up when all scores are equal
    checkDuplicateTerms(new MultiTermQuery.TopTermsScoringBooleanQueryRewrite(1024));
    checkDuplicateTerms(new MultiTermQuery.TopTermsBoostOnlyBooleanQueryRewrite(1024));
  }

  private void checkBooleanQueryBoosts(BooleanQuery bq) {
    for (BooleanClause clause : bq.clauses()) {
      final BoostQuery boostQ = (BoostQuery) clause.query();
      final TermQuery mtq = (TermQuery) boostQ.getQuery();
      assertEquals(
          "Parallel sorting of boosts in rewrite mode broken",
          Float.parseFloat(mtq.getTerm().text()),
          boostQ.getBoost(),
          0);
    }
  }

  private void checkBoosts(MultiTermQuery.RewriteMethod method) throws Exception {
    final MultiTermQuery mtq =
        new MultiTermQuery("data", method) {
          @Override
          protected TermsEnum getTermsEnum(Terms terms, AttributeSource atts) throws IOException {
            return new FilteredTermsEnum(terms.iterator()) {

              final BoostAttribute boostAtt = attributes().addAttribute(BoostAttribute.class);

              @Override
              protected AcceptStatus accept(BytesRef term) {
                boostAtt.setBoost(Float.parseFloat(term.utf8ToString()));
                if (term.length == 0) {
                  return AcceptStatus.NO;
                }
                char c = (char) (term.bytes[term.offset] & 0xff);
                if (c >= '2') {
                  if (c <= '7') {
                    return AcceptStatus.YES;
                  } else {
                    return AcceptStatus.END;
                  }
                } else {
                  return AcceptStatus.NO;
                }
              }
            };
          }

          @Override
          public String toString(String field) {
            return "dummy";
          }

          @Override
          public void visit(QueryVisitor visitor) {}
        };
    final Query q1 = searcher.rewrite(mtq);
    final Query q2 = multiSearcher.rewrite(mtq);
    final Query q3 = multiSearcherDupls.rewrite(mtq);
    if (VERBOSE) {
      System.out.println();
      System.out.println("single segment: " + q1);
      System.out.println("multi segment: " + q2);
      System.out.println("multi segment with duplicates: " + q3);
    }
    assertEquals("The multi-segment case must produce same rewritten query", q1, q2);
    assertEquals(
        "The multi-segment case with duplicates must produce same rewritten query", q1, q3);
    if (q1 instanceof MatchNoDocsQuery) {
      assertTrue(q1 instanceof MatchNoDocsQuery);
      assertTrue(q3 instanceof MatchNoDocsQuery);
    } else {
      checkBooleanQueryBoosts((BooleanQuery) q1);
      checkBooleanQueryBoosts((BooleanQuery) q2);
      checkBooleanQueryBoosts((BooleanQuery) q3);
      assert false;
    }
  }

  public void testBoosts() throws Exception {
    checkBoosts(MultiTermQuery.SCORING_BOOLEAN_REWRITE);

    // use a large PQ here to only test boosts and dont mix up when all scores are equal
    checkBoosts(new MultiTermQuery.TopTermsScoringBooleanQueryRewrite(1024));
  }

  private void checkMaxClauseLimitation(MultiTermQuery.RewriteMethod method) throws Exception {
    int savedMaxClauseCount = IndexSearcher.getMaxClauseCount();
    IndexSearcher.setMaxClauseCount(3);

    final MultiTermQuery mtq = TermRangeQuery.newStringRange("data", "2", "7", true, true, method);
    try {
      IndexSearcher.TooManyClauses expected =
          expectThrows(
              IndexSearcher.TooManyClauses.class,
              () -> {
                multiSearcherDupls.rewrite(mtq);
              });
      //  Maybe remove this assert in later versions, when internal API changes:
      assertEquals(
          "Should throw BooleanQuery.TooManyClauses with a stacktrace containing checkMaxClauseCount()",
          "checkMaxClauseCount",
          expected.getStackTrace()[0].getMethodName());
    } finally {
      IndexSearcher.setMaxClauseCount(savedMaxClauseCount);
    }
  }

  private void checkNoMaxClauseLimitation(MultiTermQuery.RewriteMethod method) throws Exception {
    int savedMaxClauseCount = IndexSearcher.getMaxClauseCount();
    IndexSearcher.setMaxClauseCount(3);

    final MultiTermQuery mtq = TermRangeQuery.newStringRange("data", "2", "7", true, true, method);
    try {
      multiSearcherDupls.rewrite(mtq);
    } finally {
      IndexSearcher.setMaxClauseCount(savedMaxClauseCount);
    }
  }

  public void testMaxClauseLimitations() throws Exception {
    checkMaxClauseLimitation(MultiTermQuery.SCORING_BOOLEAN_REWRITE);
    checkMaxClauseLimitation(MultiTermQuery.CONSTANT_SCORE_BOOLEAN_REWRITE);

    checkNoMaxClauseLimitation(MultiTermQuery.CONSTANT_SCORE_REWRITE);
    checkNoMaxClauseLimitation(MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE);
    checkNoMaxClauseLimitation(new MultiTermQuery.TopTermsScoringBooleanQueryRewrite(1024));
    checkNoMaxClauseLimitation(new MultiTermQuery.TopTermsBoostOnlyBooleanQueryRewrite(1024));
  }
}
